{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp problem_types.utils\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem Type Utils\n",
    "\n",
    "Utils to create problem types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from typing import Dict, Tuple\n",
    "\n",
    "import tensorflow as tf\n",
    "from m3tl.base_params import BaseParams\n",
    "\n",
    "\n",
    "def empty_tensor_handling_loss(labels, logits, loss_fn):\n",
    "    if tf.equal(tf.size(labels), 0):\n",
    "        return 0.0\n",
    "    if tf.equal(tf.size(tf.shape(labels)), 0):\n",
    "        return 0.0\n",
    "    if tf.equal(tf.shape(labels)[0], 0):\n",
    "        return 0.0\n",
    "    else:\n",
    "        return tf.reduce_mean(loss_fn(\n",
    "            labels, logits, from_logits=True))\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def nan_loss_handling(loss):\n",
    "    if tf.math.is_nan(loss):\n",
    "        return 0.0\n",
    "    else:\n",
    "        return loss\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def create_dummy_if_empty(inp_tensor: tf.Tensor) -> tf.Tensor:\n",
    "    shape_tensor = tf.shape(inp_tensor)\n",
    "    if tf.equal(shape_tensor[0], 0):\n",
    "        data_type = inp_tensor.dtype\n",
    "        dummy_shape_first_dim = tf.convert_to_tensor([1], dtype=tf.int32)\n",
    "        dummy_shape = tf.concat(\n",
    "            [dummy_shape_first_dim, shape_tensor[1:]], axis=0)\n",
    "        dummy_tensor = tf.zeros(dummy_shape, dtype=data_type)\n",
    "        return dummy_tensor\n",
    "    else:\n",
    "        return inp_tensor\n",
    "\n",
    "\n",
    "class BaseTop(tf.keras.Model):\n",
    "    def __init__(self, params: BaseParams, problem_name: str) -> None:\n",
    "        super(BaseTop, self).__init__(name=problem_name)\n",
    "        self.params = params\n",
    "        self.problem_name = problem_name\n",
    "\n",
    "    def call(self, inputs: Tuple[Dict], mode: str):\n",
    "        raise NotImplementedError\n",
    "    \n",
    "def pad_to_shape(from_tensor: tf.Tensor, to_tensor: tf.Tensor, axis=1) -> tf.Tensor:\n",
    "    # sometimes the length of labels dose not equal to length of inputs\n",
    "    # that's caused by tf.data.experimental.bucket_by_sequence_length in multi problem scenario\n",
    "    pad_len = tf.shape(input=to_tensor)[\n",
    "        axis] - tf.shape(input=from_tensor)[axis]\n",
    "\n",
    "    # top, bottom, left, right\n",
    "    pad_tensor = [[0, 0] for _ in range(len(from_tensor.shape))]\n",
    "    pad_tensor[axis] = [0, pad_len]\n",
    "    from_tensor = tf.pad(tensor=from_tensor, paddings=pad_tensor)\n",
    "    return from_tensor\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
